from llm import llm, llm0,llm0_1, llm_fc
from tools import tools_calc, tools_say, exec_call

import inspect

from langgraph.prebuilt import ToolExecutor

tool_executor = ToolExecutor(tools_calc)


llm_calc = llm_fc.bind_tools(tools=tools_calc)


import operator
from typing import Annotated, Sequence, TypedDict
from langchain_core.messages import BaseMessage

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]


from langchain_core.messages import ToolMessage
from langgraph.prebuilt import ToolInvocation

# Define the function that determines whether to continue or not
def should_continue(state):
    messages = state["messages"]
    last_message = messages[-1]
    # If there is no function call, then we finish
    if not last_message.tool_calls:
        return "end"
    # Otherwise if there is, we continue
    else:
        return "continue"

def prn_call(name,state):
    return
    print(name, ': ')
    for msg in state["messages"]:
        if msg.content == '':
            print(
                '       calls: ',
                msg.tool_calls[0]['name'],
                msg.tool_calls[0]['args'],
            )
        else:
            print('    message: ',msg.content)

def get_fn_name():
    return inspect.currentframe().f_back.f_code.co_name 


# Define the function that calls the model
def call_model(state):
    prn_call(get_fn_name(), state)
    messages = state["messages"]
    response = llm_calc.invoke(messages)
    # We return a list, because this will get added to the existing list
    return {"messages": [response]}


# Define the function to execute tools
def call_tool(state):
    prn_call(get_fn_name(), state)
    messages = state["messages"]
    # Based on the continue condition
    # we know the last message involves a function call
    last_message = messages[-1]
    # We construct an ToolInvocation from the function_call
    tool_call = last_message.tool_calls[0]
    action = ToolInvocation(
        tool=tool_call["name"],
        tool_input=tool_call["args"],
    )
    # We call the tool_executor and get back a response
    response = tool_executor.invoke(action)
    # We use the response to create a FunctionMessage
    function_message = ToolMessage(
        content=str(response), name=action.tool, tool_call_id=tool_call["id"]
    )
    # We return a list, because this will get added to the existing list
    return {"messages": [function_message]}


from langgraph.graph import END, StateGraph

# Define a new graph
workflow = StateGraph(AgentState)

# Define the two nodes we will cycle between
workflow.add_node("agent", call_model)
workflow.add_node("action", call_tool)

# Set the entrypoint as `agent`
# This means that this node is the first one called
workflow.set_entry_point("agent")

# We now add a conditional edge
workflow.add_conditional_edges(
    # First, we define the start node. We use `agent`.
    # This means these are the edges taken after the `agent` node is called.
    "agent",
    # Next, we pass in the function that will determine which node is called next.
    should_continue,
    # Finally we pass in a mapping.
    # The keys are strings, and the values are other nodes.
    # END is a special node marking that the graph should finish.
    # What will happen is we will call `should_continue`, and then the output of that
    # will be matched against the keys in this mapping.
    # Based on which one it matches, that node will then be called.
    {
        # If `tools`, then we call the tool node.
        "continue": "action",
        # Otherwise we finish.
        "end": END,
    },
)

# We now add a normal edge from `tools` to `agent`.
# This means that after `tools` is called, `agent` node is called next.
workflow.add_edge("action", "agent")

# Finally, we compile it!
# This compiles it into a LangChain Runnable,
# meaning you can use it as you would any other runnable
app = workflow.compile()


app.get_graph().print_ascii()



from langchain_core.messages import HumanMessage

q = '7乘以8等于'

def invoke():
    ret = app.invoke(
        {"messages": [HumanMessage(content=q)]}
    )

    print(ret['messages'][-1])

def stream():
    inputs = {"messages": [HumanMessage(content=q)]}
    for output in app.stream(inputs):
        # stream() yields dictionaries with output keyed by node name
        for key, value in output.items():
            print(f"Output from node '{key}':")
            print("---")
            print(value)
        print("\n---\n")


# invoke()
stream()
